import torch
import torch.nn as nn

from .umetrack_model import UmeTrackModel
from .model_input import InputFrameData, InputFrameDesc, InputSkeletonData


class UmeTrackWrapper(nn.Module):
    def __init__(self, umetrack: UmeTrackModel):
        super().__init__()
        self.model = umetrack

    def forward(self, images: torch.Tensor, extrinsics_xf: torch.Tensor,
                intrinsics: torch.Tensor, joint_axes: torch.Tensor,
                joint_rest: torch.Tensor, sample_range: torch.Tensor,
                memory_idx: torch.Tensor, use_memory: torch.Tensor,
                hand_idx: torch.Tensor):
        # 将原本的 InputFrameData、Desc、SkeletonData 转成张量输入
        frame_data = InputFrameData(
            left_images=images,
            extrinsics_xf=extrinsics_xf,
            intrinsics=intrinsics,
        )
        frame_desc = InputFrameDesc(
            sample_range=sample_range,
            memory_idx=memory_idx,
            use_memory=use_memory,
            hand_idx=hand_idx,
        )
        skel_data = InputSkeletonData(
            joint_rotation_axes=joint_axes,
            joint_rest_positions=joint_rest,
        )
        # 调用你的回归接口
        out = self.model.regress_pose_use_skeleton(frame_data, frame_desc, skel_data)
        # 比如只返回 wrist_xfs; 如果要更多输出，做个 tuple
        return out.wrist_xfs
